#!/usr/bin/env python3
"""
BaoStock数据源工具
提供统一的BaoStock数据访问接口
"""

import baostock as bs
import pandas as pd
from typing import Dict, Optional, Any
import logging

# 导入日志模块
from tradingagents.utils.logging_manager import get_logger
logger = get_logger('agents')


class BaoStockProvider:
    """BaoStock数据提供者"""
    
    def __init__(self):
        """初始化BaoStock提供者"""
        self.lg = None
        self._login()
    
    def _login(self):
        """登录BaoStock"""
        try:
            self.lg = bs.login()
            if self.lg.error_code != '0':
                logger.error(f"❌ BaoStock登录失败: {self.lg.error_msg}")
                return False
            logger.info("✅ BaoStock登录成功")
            return True
        except Exception as e:
            logger.error(f"❌ BaoStock登录异常: {e}")
            return False
    
    def _logout(self):
        """登出BaoStock"""
        try:
            if self.lg and self.lg.error_code == '0':
                bs.logout()
                logger.info("✅ BaoStock登出成功")
        except Exception as e:
            logger.error(f"❌ BaoStock登出异常: {e}")
    
    def get_stock_data(self, symbol: str, start_date: str, end_date: str) -> Optional[pd.DataFrame]:
        """
        获取股票历史数据
        
        Args:
            symbol: 股票代码
            start_date: 开始日期 (YYYY-MM-DD)
            end_date: 结束日期 (YYYY-MM-DD)
            
        Returns:
            DataFrame: 股票数据
        """
        try:
            # 转换股票代码格式
            if symbol.startswith('6'):
                bs_code = f"sh.{symbol}"
            else:
                bs_code = f"sz.{symbol}"
            
            # 查询历史行情数据
            rs = bs.query_history_k_data_plus(
                code=bs_code,
                fields="date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST",
                start_date=start_date,
                end_date=end_date,
                frequency="d",  # 日线
                adjustflag="3"  # 复权类型: 3-不复权
            )
            
            if rs.error_code != '0':
                logger.error(f"❌ BaoStock数据查询失败: {rs.error_msg}")
                return None
            
            # 获取数据
            data_list = []
            while (rs.error_code == '0') & rs.next():
                data_list.append(rs.get_row_data())
            
            if not data_list:
                logger.warning(f"⚠️ BaoStock未获取到{symbol}的数据")
                return None
            
            # 转换为DataFrame
            columns = ["date", "code", "open", "high", "low", "close", "preclose", "volume", "amount", 
                      "adjustflag", "turn", "tradestatus", "pctChg", "isST"]
            df = pd.DataFrame(data_list, columns=columns)
            
            # 数据类型转换
            numeric_columns = ["open", "high", "low", "close", "preclose", "volume", "amount", "turn", "pctChg"]
            for col in numeric_columns:
                df[col] = pd.to_numeric(df[col], errors='coerce')
            
            # 日期转换
            df['date'] = pd.to_datetime(df['date'])
            
            logger.info(f"✅ BaoStock成功获取{symbol}数据，共{len(df)}条记录")
            return df
            
        except Exception as e:
            logger.error(f"❌ BaoStock获取{symbol}数据异常: {e}")
            return None
    
    def get_stock_info(self, symbol: str) -> Dict[str, Any]:
        """
        获取股票基本信息
        
        Args:
            symbol: 股票代码
            
        Returns:
            Dict: 股票基本信息
        """
        try:
            # 转换股票代码格式
            if symbol.startswith('6'):
                bs_code = f"sh.{symbol}"
            else:
                bs_code = f"sz.{symbol}"
            
            # 查询股票基本信息
            rs = bs.query_stock_basic(code=bs_code)
            if rs.error_code != '0':
                logger.error(f"❌ BaoStock查询股票信息失败: {rs.error_msg}")
                return {'symbol': symbol, 'name': f'股票{symbol}', 'source': 'baostock'}
            
            # 获取数据
            data_list = []
            while (rs.error_code == '0') & rs.next():
                data_list.append(rs.get_row_data())
            
            if data_list:
                # BaoStock返回格式: [code, code_name, ipoDate, outDate, type, status]
                info = {
                    'symbol': symbol,
                    'name': data_list[0][1],  # code_name
                    'source': 'baostock',
                    'list_date': data_list[0][2],  # ipoDate
                    'type': data_list[0][4],  # type
                    'status': data_list[0][5]  # status
                }
                logger.info(f"✅ BaoStock成功获取{symbol}基本信息")
                return info
            else:
                logger.warning(f"⚠️ BaoStock未获取到{symbol}的基本信息")
                return {'symbol': symbol, 'name': f'股票{symbol}', 'source': 'baostock'}
                
        except Exception as e:
            logger.error(f"❌ BaoStock获取{symbol}信息异常: {e}")
            return {'symbol': symbol, 'name': f'股票{symbol}', 'source': 'baostock', 'error': str(e)}


# 全局实例
_baostock_provider = None


def get_baostock_provider() -> BaoStockProvider:
    """获取BaoStock提供者实例"""
    global _baostock_provider
    if _baostock_provider is None:
        _baostock_provider = BaoStockProvider()
    return _baostock_provider


def test_baostock_connection():
    """测试BaoStock连接"""
    try:
        provider = get_baostock_provider()
        if provider.lg and provider.lg.error_code == '0':
            print("✅ BaoStock连接测试成功")
            return True
        else:
            print("❌ BaoStock连接测试失败")
            return False
    except Exception as e:
        print(f"❌ BaoStock连接测试异常: {e}")
        return False


if __name__ == "__main__":
    test_baostock_connection()